# 这段代码实现了一个基于 LangChain 框架的 AI 代理（Agent），它能够通过 MCP 协议与外部工具服务器通信 - 智能代理
# 此文件实现了一个基于 LangChain 框架的 AI 代理，该代理借助 MCP 协议与外部工具服务器通信。在客户端 - 服务器架构里，
# 该文件扮演客户端的角色，负责与服务器建立连接、获取工具并调用 AI 代理来处理问题。
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client ## 导入 MCP 标准输入输出客户端
from langchain_mcp_adapters.tools import load_mcp_tools #从 LangChain MCP 适配器中加载 MCP 工具
from langgraph.prebuilt import create_react_agent #从 LangGraph 预构建模块中导入创建 ReAct 代理的函数
from langchain_ollama import ChatOllama # 从 LangChain Ollama 模块中导入聊天模型类
import asyncio # 导入异步 I/O 模块

# 初始化一个基于 Ollama 的聊天模型，使用 qwen2.5:72b 模型
model = ChatOllama(model='qwen2.5:72b')

# 定义标准输入输出服务器的参数 -MCP通信，连接到本地的 math_server.py 文件
server_params = StdioServerParameters(
    command="python",
    # Make sure to update to the full absolute path to your math_server.py file
    args=["math_server.py"],
)

#  生成一个智能体，必须在异步环境下连接MCP服务端
async def run_agent():
    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:
            # Initialize the connection
            await session.initialize()

            # Get tools
            tools = await load_mcp_tools(session)

            # Create and run the agent
            agent = create_react_agent(model, tools)
            agent_response = await agent.ainvoke({"messages": "what's (3 + 5) x 12?"})
            return agent_response

# Run the async function
if __name__ == "__main__":
    result = asyncio.run(run_agent())
    print(result)